import numpy as np
import argparse
import sys
import pandas as pd
import subprocess as sp
from pathlib import Path
from matplotlib import pyplot as plt
from atm_mbar.base import read_config, read_yaml, read_u_energy, bias_fcn
from atm_mbar.analysis import pre_data, fep, time_convergence, block_average

parser = argparse.ArgumentParser()
parser.add_argument('-n', '--name', required=True, help='name')
parser.add_argument('-c', '--config', required=True, help='config file')
parser.add_argument('-s', '--start', required=False,
                    help='start sample', default=0)
parser.add_argument('-e', '--end', required=False,
                    help='end sample', default=sys.maxsize)

args = parser.parse_args()

def main():
    config = dict()
    try:
        config = read_yaml(args.config)
    except Exception:
        try:
            config = read_config(args.config)
        except Exception:
            raise RuntimeError('can not read config.')

    lambda1 = np.array(config['LAMBDA1'], dtype=np.float64)
    lambda2 = np.array(config['LAMBDA2'], dtype=np.float64)
    lambdas = np.array(config['LAMBDAS'], dtype=np.float64)

    alpha = np.array(config['ALPHA'], dtype=np.float64)
    u0 = np.array(config['U0'], dtype=np.float64)
    w0 = np.array(config['W0COEFF'], dtype=np.float64)
    (lig1_usamples, lig2_usamples, N_k1, N_k2) = read_u_energy(
        args.name, len(lambdas), int(args.start), int(args.end))
    nstates = len(lambdas) // 2
    u_kn1 = np.zeros((nstates, len(lig1_usamples)))
    u_kn2 = np.zeros((nstates, len(lig2_usamples)))

    for k in range(nstates):
        u_kn1[k] = bias_fcn(lig1_usamples, 1. / (1.986e-3 * 300.0),
                            lambda1[k], lambda2[k], alpha[k], u0[k], w0[k])
    for k in range(nstates, len(lambdas)):
        u_kn2[k - nstates] = bias_fcn(lig2_usamples, 1. / (1.986e-3 * 300.0),
                                      lambda1[k], lambda2[k], alpha[k], u0[k], w0[k])
    u_nk1 = pd.DataFrame(u_kn1)
    u_nk2 = pd.DataFrame(u_kn2)


    u_nk1, u_nk2 = pre_data(u_nk1, u_nk2, n_k1=N_k1, n_k2=N_k2,lambdas=lambdas)
    mbar_result = fep(u_nk1, u_nk2)
    with open(f'mbar_{args.name}.result' , 'w') as f:
        f.write(str(mbar_result))
    time_convergence(u_nk1, u_nk2)
    print('--------------------------------')
    block_average(u_nk1, u_nk2)
    print('--------------------------------')
    path = Path(__file__).parent
    cmd = ['Rscript',str(path.joinpath("calc.R")), args.config, args.name, str(args.start), str(999999)]
    sp.run(cmd)
    lambda1_data = {}
    lambda2_data = {}

    for lambda1 in Path().glob('lambda1-*.dat'):
        lambda1_data[lambda1.stem[lambda1.stem.find(
            '-') + 1:]] = np.loadtxt(lambda1)
    for lambda2 in Path().glob('lambda2-*.dat'):
        lambda2_data[lambda2.stem[lambda2.stem.find(
            '-') + 1:]] = np.loadtxt(lambda2)
    DGb = Path(f'{args.name}.result').read_text()
    plt.figure()
    for l,i in lambda1_data.items():
        plt.plot(i[:,0], i[:,1], label=f'lambda {l}')
    plt.xlabel('X')
    plt.ylabel('Y')
    plt.title('lambda 1')
    plt.legend()

    plt.savefig('lambda1.png')
    plt.close()
    plt.figure()
    for l,i in lambda2_data.items():
        plt.plot(i[:,0], i[:,1], label=f'lambda {l}')
    plt.xlabel('X')
    plt.ylabel('Y')
    plt.title('lambda 2')
    plt.legend()

    plt.savefig('lambda2.png')
    plt.close()
    print(DGb)

if __name__ == "__main__":
    pass
